package net.cyue.ort;

import ai.onnxruntime.OrtEnvironment;
import ai.onnxruntime.OrtException;
import ai.onnxruntime.OrtSession;

import net.cyue.ort.util.ONNXRuntimeBackendUtil;

import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.List;
import java.util.Objects;
import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;

/**
 * 负责管理 OrtEnvironment、OrtSession 以及 ByteBuffer 缓冲池
 * 将底层运行时资源从 LLM 主类中抽离，提升可维护性
 */
public class ModelSessionManager implements AutoCloseable {

    private final OrtEnvironment environment;
    private final OrtSession session;
    private final Queue<ByteBuffer> bufferPool = new ConcurrentLinkedQueue<>();
    private static final int DEFAULT_BUFFER_CAPACITY = 4 * 1024 * 1024; // 4MB

    public ModelSessionManager(List<ByteBuffer> modelBuffers) {
        Objects.requireNonNull(modelBuffers, "modelBuffers cannot be null");
        if (modelBuffers.isEmpty()) {
            throw new IllegalArgumentException("modelBuffers cannot be empty");
        }

        try {
            this.environment = OrtEnvironment.getEnvironment();
            OrtSession.SessionOptions options = createSessionOptions();
            byte[] modelData = modelBuffers.size() == 1
                ? readBytes(modelBuffers.get(0))
                : mergeBuffers(modelBuffers);

            this.session = environment.createSession(modelData, options);
            initializeBufferPool();

            // 主动触发一次 GC，确保模型加载后的 DirectByteBuffer 及时回收
            System.gc();
        } catch (OrtException e) {
            throw new RuntimeException("Failed to initialize ONNX session", e);
        }
    }

    private OrtSession.SessionOptions createSessionOptions() throws OrtException {
        OrtSession.SessionOptions options = new OrtSession.SessionOptions();
        options.setIntraOpNumThreads(Runtime.getRuntime().availableProcessors());
        options.setOptimizationLevel(OrtSession.SessionOptions.OptLevel.ALL_OPT);
        options.setMemoryPatternOptimization(true);
        options.addCPU(true);

        if (ONNXRuntimeBackendUtil.isSupportDirectML()) {
            System.out.println("使用 DirectML");
            options.addDirectML(0);
        } else {
            System.out.println("DirectML is not supported");
        }

        if (ONNXRuntimeBackendUtil.isSupportCUDA()) {
            System.out.println("使用 CUDA");
            options.addCUDA();
        } else {
            System.out.println("CUDA is not supported");
        }

        return options;
    }

    private void initializeBufferPool() {
        int processors = Math.max(1, Runtime.getRuntime().availableProcessors());
        for (int i = 0; i < processors; i++) {
            bufferPool.offer(
                ByteBuffer.allocateDirect(DEFAULT_BUFFER_CAPACITY)
                    .order(ByteOrder.LITTLE_ENDIAN)
            );
        }
    }

    private byte[] readBytes(ByteBuffer buffer) {
        byte[] modelData = new byte[buffer.remaining()];
        buffer.get(modelData);
        return modelData;
    }

    private byte[] mergeBuffers(List<ByteBuffer> buffers) {
        int totalSize = buffers.stream()
            .mapToInt(ByteBuffer::remaining)
            .sum();

        ByteBuffer merged = ByteBuffer.allocate(totalSize);
        buffers.forEach(buffer -> {
            byte[] data = new byte[buffer.remaining()];
            buffer.get(data);
            merged.put(data);
        });

        return merged.array();
    }

    public OrtSession getSession() {
        return session;
    }

    public OrtEnvironment getEnvironment() {
        return environment;
    }

    public ByteBuffer borrowBuffer(int size) {
        ByteBuffer buffer = bufferPool.poll();
        if (buffer != null && buffer.capacity() >= size) {
            buffer.clear();
            return buffer;
        }
        return ByteBuffer.allocateDirect(size).order(ByteOrder.LITTLE_ENDIAN);
    }

    public void recycleBuffer(ByteBuffer buffer) {
        if (buffer != null) {
            buffer.clear();
            bufferPool.offer(buffer);
        }
    }

    @Override
    public void close() {
        try {
            if (session != null) {
                session.close();
            }
            if (environment != null) {
                environment.close();
            }
            bufferPool.forEach(ByteBuffer::clear);
            bufferPool.clear();
        } catch (Exception e) {
            throw new RuntimeException("Failed to close ONNX runtime resources", e);
        }
    }
}

